#!/usr/bin/env python
from os.path import exists,isfile
from os import getcwd,chdir,mkdir,listdir,system
import numpy as np
from ase.io import read
''' 
Setting env:
    export LD_LIBRARY_PATH=/home/feng/mlff/cc/shengbte/spglib/lib/:$LD_LIBRARY_PATH 
'''

def write_control(nx=1,ny=1,nz=1,t=300,sx=8,sy=8,sz=1):
    atoms = read('POSCAR.unitcell')
    symbols = atoms.get_chemical_symbols()
    elements = list(set(symbols))

    cell = atoms.get_cell()
    cell = cell[:].astype(dtype=np.float32)
    rcell     = np.linalg.inv(cell).astype(dtype=np.float32)
    positions = atoms.get_positions()
    xf        = np.dot(positions,rcell)
    xf        = np.mod(xf,1.0)
    with open('CONTROL','w') as f:
         f.write('&allocations\n')
         f.write('        nelements=     {:d},\n'.format(len(elements)))
         f.write('        natoms=        {:d},\n'.format(len(atoms)))
         f.write('        ngrid(:)=      {:d} {:d} {:d} \n'.format(nx,ny,nz))
         f.write('&end\n')
         f.write('&crystal\n')
         f.write('        lfactor=0.1,\n')
         f.write('        lattvec(:,1)= {:f}  {:f}  {:f},\n'.format(cell[0][0],cell[0][1],cell[0][2]))
         f.write('        lattvec(:,2)= {:f}  {:f}  {:f},\n'.format(cell[1][0],cell[1][1],cell[1][2]))
         f.write('        lattvec(:,3)= {:f}  {:f}  {:f},\n'.format(cell[2][0],cell[2][1],cell[2][2]))
         f.write('        elements= \"{:s}\"\n'.format(' '.join(elements)))
         types = [str(elements.index(symbols[i])+1) for i in range(len(atoms))]
         f.write('        types=   {:s},\n'.format(' '.join(types)))
         
         for i in range(len(atoms)):
             f.write('        positions(:,{:d})= {:f} {:f} {:f},\n'.format(i+1,
                                        xf[i][0],xf[i][1],xf[i][2]))
         f.write('        epsilon(:,1)= 1.0 0.0 0.0, \n')
         f.write('        epsilon(:,2)= 0.0 1.0 0.0, \n')
         f.write('        epsilon(:,3)= 0.0 0.0 1.0, \n')
         for i in range(len(atoms)):
             f.write('        born(:,1,{:d})= 0.0 0.0 0.0,\n'.format(i+1))
             f.write('        born(:,2,{:d})= 0.0 0.0 0.0,\n'.format(i+1))
             f.write('        born(:,3,{:d})= 0.0 0.0 0.0,\n'.format(i+1))
         f.write('        scell(:)= {:3d}  {:3d} {:3d} \n'.format(sx,sy,sz))
         f.write('&end\n')
         f.write('&parameters\n')
         f.write('        T= {:.1f} \n'.format(t))
         f.write('&end\n')
         f.write('&flags\n')
         f.write('        convergence=.True. \n')
         f.write('        nonanalytic=.True.  \n')
         f.write('&end\n')

rwd      = getcwd()                    # root work directory

with open('kappa_vs_qgrid.txt','w') as f:
     print('# nx ny nz  kappa',file=f)

z = 1
for x in (18,20): # range(10,15,4):
    y = x
    wd       = 'N{:d}{:d}{:d}'.format(x,y,z)
    chdir(rwd)
    if not exists(wd):
       mkdir(wd)
    chdir(rwd+'/'+wd)
    system('cp ../ShengBTE ShengBTE')
    system('cp ../POSCAR.unitcell POSCAR.unitcell')
    system('cp ../FORCE_CONSTANTS_2ND FORCE_CONSTANTS_2ND')
    system('cp ../FORCE_CONSTANTS_3RD FORCE_CONSTANTS_3RD')
    write_control(nx=x,ny=y,nz=z)
    system('mpirun -n 16 ./ShengBTE >sheng.out')

    with open('T300K/BTE.kappa_scalar','r') as f:
         lines = f.readlines()
         l = lines[-1].split()
         if len(l)==0:
            l = lines[-2].split()
         else:
            l = lines[-1].split()
            
    chdir(rwd)

    if l[1] !='NaN':
       try:
          k = float(l[1])
       except:
          continue
       if k>0.0 and k<=10000:
          with open('kappa_vs_qgrid.txt','a') as f:
                print('{:3d} {:3d} {:3d}  {:11.6f}'.format(x,y,z,k),file=f)
